// STL
#include <fstream> // std::ifstream, std::ofstream
#include <string>  // std::string
#include <vector>  // std::vector<>

// Boost
#include <boost/archive/text_iarchive.hpp> // boost::archive::text_iarchive
#include <boost/archive/text_oarchive.hpp> // boost::archive::text_oarchive

// stats++
#include "statsxx/machine_learning/DBN.hpp" // machine_learning::DBN


//
// DESC: Creates a DBN (or reads one in).
//
machine_learning::DBN create_DBN(
                                 const bool              create,
                                 // -----
                                 const std::vector<int> &create_nunits,
                                 const std::vector<int> &create_units_type,
                                 // -----
                                 const std::string       dbn_file
                                 )
{
    machine_learning::DBN dbn;
    
    if(create)
    {
        int nRBM = create_nunits.size()-1;
        
        std::vector<int> nv(nRBM);
        std::vector<int> nh(nRBM);
        // ---
        std::vector<int> vtype(nRBM);
        std::vector<int> htype(nRBM);
        
        for(int i = 0; i < nRBM; ++i)
        {  
            nv[i]    = create_nunits[i];
            nh[i]    = create_nunits[i+1];
            
            vtype[i] = create_units_type[i];
            htype[i] = create_units_type[i+1];
        }
        
        dbn = machine_learning::DBN(
                                    nRBM,
                                    // ---
                                    nv,
                                    nh,
                                    // ---
                                    vtype,
                                    htype
                                    );
        
        // SAVE DBN TO ARCHIVE
        {
            std::ofstream ofs(dbn_file);
            
            boost::archive::text_oarchive oa(ofs);
            
            oa << dbn;
        }
    }
    else
    {
        // READ DBN FROM ARCHIVE
        {
            std::ifstream ifs(dbn_file);
            
            boost::archive::text_iarchive ia(ifs);
            
            ia >> dbn;
        }
    }
    
    return dbn;    
}
